import numpy as np
from points import Point

class PointsHandler:
    def __init__(self):
        self.points = []

    def add_point(self, point):
        self.points.append(point)

    def apply_rf_pulse(self):
        for point in self.points:
            point.clear()

    def apply_y_gradient(self,x_gradient, y_gradient, duration):
        for point in self.points:
            point.apply_y_gradient(x_gradient=x_gradient,y_gradient=y_gradient, duration=duration)

    def read_out(self, x_gradient, duration, sampling_rate):
        signals = []
        t = None
        for point in self.points:
            t, signal = point.read_out(x_gradient, duration, sampling_rate)
            signals.append(signal)
        summed_signal = np.sum(signals, axis=0)  # 返回所有点的信号求和
        return t, summed_signal

    def get_point(self, position):
        for point in self.points:
            if np.array_equal(point.position, position):
                return point
        return None

    def get_all_points(self):
        return self.points

# 单元测试部分
if __name__ == "__main__":
    import matplotlib.pyplot as plt

    def test_points_handler():
        # 创建 PointsHandler 对象
        handler = PointsHandler()

        # 添加多个 Point 对象
        handler.add_point(Point([1, 1], 1, T2=0.03, B0=1000.0))
        handler.add_point(Point([3, 3], 1, T2=0.01, B0=1000.0))
        # handler.add_point(Point([1, 5], 1, T2=0.02, B0=1000.0))
        # handler.add_point(Point([5, 1], 1, T2=0.03, B0=1000.0))

        # 应用 RF 脉冲
        handler.apply_rf_pulse()

        # 应用 y 方向梯度
        handler.apply_y_gradient(x_gradient=-100.0,y_gradient=100, duration=0.1)

        # 读取 x 方向梯度信号
        t, summed_signal = handler.read_out(x_gradient=-1000.0, duration=0.2, sampling_rate=1000)

        # 画出信号的实部和虚部
        plt.figure(figsize=(6, 6))

        plt.subplot(2, 1, 1)
        plt.plot(t, summed_signal.real)
        plt.title('Real Part of Summed Signal')
        plt.xlabel('Time (s)')
        plt.ylabel('Amplitude')

        plt.subplot(2, 1, 2)
        plt.plot(t, np.abs(summed_signal))
        plt.title('Magnitude of Summed Signal')
        plt.xlabel('Time (s)')
        plt.ylabel('Amplitude')

        plt.tight_layout()
        plt.show()

    test_points_handler()